"""
生产环境的Vanna AI数据查询服务
优化版本，包含配置管理、错误处理、日志记录等功能
"""

import os
import logging
from logging.handlers import RotatingFileHandler, TimedRotatingFileHandler
from typing import Optional, Dict, Any
from dataclasses import dataclass
from functools import wraps
import asyncio
from datetime import datetime
import pathlib
import nest_asyncio
import uvicorn
from typing import List

# 应用nest_asyncio以允许嵌套事件循环
nest_asyncio.apply()

# 导入Vanna相关模块
from vanna.integrations.ollama import OllamaLlmService
from vanna.tools import RunSqlTool
from vanna.integrations.mysql import MySQLRunner
from vanna.core.tool import ToolContext, ToolResult
from vanna.tools.agent_memory import SaveQuestionToolArgsTool, SearchSavedCorrectToolUsesTool, SaveTextMemoryTool
from vanna.integrations.chromadb import ChromaAgentMemory
from vanna.core.user import UserResolver, User, RequestContext
from vanna import Agent
from vanna.core.registry import ToolRegistry
from vanna.tools import VisualizeDataTool
from vanna.core.agent import AgentConfig
from vanna.servers.fastapi import VannaFastAPIServer
from vanna.core.enhancer import LlmContextEnhancer
from vanna.core.llm import LlmMessage
from vanna.capabilities.sql_runner import RunSqlToolArgs

# 配置日志
def setup_logging():
    """设置日志配置，包括控制台和文件输出"""
    # 创建logs目录（如果不存在）
    logs_dir = pathlib.Path("logs")
    logs_dir.mkdir(exist_ok=True)
    
    # 获取当前日期作为日志文件名的一部分
    current_date = datetime.now().strftime("%Y%m%d")
    log_file = logs_dir / f"vanna_prod_{current_date}.log"
    
    # 配置日志格式
    formatter = logging.Formatter(
        '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )
    
    # 创建根日志记录器
    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    
    # 清除现有的处理器
    for handler in root_logger.handlers[:]:
        root_logger.removeHandler(handler)
    
    # 添加控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)
    root_logger.addHandler(console_handler)
    
    # 添加文件处理器（按大小轮转）
    file_handler = RotatingFileHandler(
        log_file,
        maxBytes=10*1024*1024,  # 10MB
        backupCount=5,
        encoding='utf-8'
    )
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)
    root_logger.addHandler(file_handler)
    
    # 添加按时间轮转的处理器（每天午夜轮转）
    time_handler = TimedRotatingFileHandler(
        logs_dir / "vanna_prod_daily.log",
        when='midnight',
        interval=1,
        backupCount=30,
        encoding='utf-8'
    )
    time_handler.setLevel(logging.INFO)
    time_handler.setFormatter(formatter)
    root_logger.addHandler(time_handler)
    
    # 记录日志系统启动信息
    logger = logging.getLogger(__name__)
    logger.info(f"日志系统已初始化，日志文件: {log_file}")

# 初始化日志系统
setup_logging()
logger = logging.getLogger(__name__)

# 常用表配置
COMMON_TABLES = ['test_employee', 'test_pay_sheet']

# 业务上下文：用于存储动态生成的建表语句
business_contexts: List[str] = []

def handle_errors(func):
    """错误处理装饰器，用于捕获和记录异常"""
    @wraps(func)
    async def wrapper(*args, **kwargs):
        try:
            return await func(*args, **kwargs)
        except Exception as e:
            logger.error(f"在执行 {func.__name__} 时发生错误: {str(e)}", exc_info=True)
            raise
    return wrapper


@dataclass
class DatabaseConfig:
    """数据库配置类"""
    host: str
    database: str
    user: str
    password: str
    port: int
    
    @classmethod
    def from_env(cls) -> 'DatabaseConfig':
        """从环境变量创建数据库配置"""
        return cls(
            host=os.getenv('DB_HOST', '192.168.1.246'),
            database=os.getenv('DB_NAME', 'test_employee'),
            user=os.getenv('DB_USER', 'root'),
            password=os.getenv('DB_PASSWORD', 'aG9uZ2Nhbg==#2023'),
            port=int(os.getenv('DB_PORT', '3318'))
        )


@dataclass
class LLMConfig:
    """LLM配置类"""
    model: str
    host: str
    
    @classmethod
    def from_env(cls) -> 'LLMConfig':
        """从环境变量创建LLM配置"""
        return cls(
            model=os.getenv('LLM_MODEL', 'gpt-oss:latest'),
            host=os.getenv('LLM_HOST', 'http://zh.daoreach.com:11434')
        )


@dataclass
class ServerConfig:
    """服务器配置类"""
    port: int
    max_tool_iterations: int
    temperature: float
    
    @classmethod
    def from_env(cls) -> 'ServerConfig':
        """从环境变量创建服务器配置"""
        return cls(
            port=int(os.getenv('SERVER_PORT', '8001')),
            max_tool_iterations=int(os.getenv('MAX_TOOL_ITERATIONS', '30')),
            temperature=float(os.getenv('LLM_TEMPERATURE', '0.7'))
        )


class ProductionEditableRunSqlTool(RunSqlTool):
    """生产环境SQL执行工具，增强错误处理和日志记录"""
    
    @handle_errors
    async def execute(self, context: ToolContext, args) -> ToolResult:
        """
        执行SQL查询，记录生成的SQL语句
        
        Args:
            context: 工具执行上下文
            args: SQL执行参数
            
        Returns:
            ToolResult: SQL执行结果
        """
        original_sql = args.sql
        logger.info(f"用户 {context.user.email} 执行SQL查询:\n{original_sql}")
        
        # 记录执行开始时间
        import time
        start_time = time.time()
        
        # 执行SQL查询
        result = await super().execute(context, args)
        
        # 记录执行时间
        execution_time = time.time() - start_time
        logger.info(f"SQL查询执行完成，耗时: {execution_time:.2f}秒")
        
        return result


class ProductionUserResolver(UserResolver):
    """生产环境用户解析器，增强安全性"""
    
    async def resolve_user(self, request_context: RequestContext) -> User:
        """
        解析用户信息，支持多种认证方式
        
        Args:
            request_context: 请求上下文
            
        Returns:
            User: 解析后的用户对象
        """
        # 尝试从Cookie获取用户邮箱
        user_email = request_context.get_cookie('vanna_email')
        
        # 如果Cookie中没有，尝试从Header获取
        if not user_email:
            user_email = request_context.get_header('X-User-Email')
        
        # 如果都没有，使用默认访客用户
        if not user_email:
            user_email = 'guest@example.com'
            logger.warning("未找到用户信息，使用默认访客账户")
        
        # 根据邮箱确定用户组
        group = 'admin' if user_email == 'admin@example.com' else 'user'
        
        logger.info(f"用户 {user_email} 已认证，组: {group}")
        return User(id=user_email, email=user_email, group_memberships=[group])


class ProductionLlmContextEnhancer(LlmContextEnhancer):
    """生产环境LLM上下文增强器，优化表结构获取和缓存"""
    
    def __init__(self, mysql_runner: MySQLRunner, agent_memory=None):
        """
        初始化LLM上下文增强器
        
        Args:
            mysql_runner: MySQL运行器实例
            agent_memory: 代理内存实例
        """
        self.mysql_runner = mysql_runner
        self.agent_memory = agent_memory
        self.table_schema_cache = {}
        self.all_tables = []
        self._initialized = False
        self._init_lock = asyncio.Lock()
    
    @handle_errors
    async def _ensure_initialized(self):
        """确保已初始化表结构信息，使用锁防止并发初始化"""
        if not self._initialized:
            async with self._init_lock:
                if not self._initialized:
                    await self._get_all_tables()
                    await self._get_common_table_structures()
                    self._initialized = True
                    logger.info(f"表结构初始化完成，共加载 {len(self.table_schema_cache)} 个表结构")
    
    @handle_errors
    async def _get_all_tables(self):
        """获取数据库中的所有表名"""
        if self.all_tables:
            return self.all_tables
            
        try:
            context = ToolContext(
                user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
                conversation_id="schema_loading",
                request_id="schema_request",
                agent_memory=self.agent_memory
            )
            
            tables_query = "SHOW TABLES"
            result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=tables_query),
                context
            )
            
            if not result.empty:
                self.all_tables = result.iloc[:, 0].tolist()
                logger.info(f"获取到数据库表列表: {self.all_tables}")
            
        except Exception as e:
            logger.error(f"获取表列表时出错: {e}")
            raise
            
        return self.all_tables
    
    @handle_errors
    async def _get_common_table_structures(self) -> None:
        """获取常用表的结构"""
        # 确保COMMON_TABLES中指定的表结构都已加载
        for table_name in COMMON_TABLES:
            # 如果表结构尚未缓存，则获取它
            if table_name not in self.table_schema_cache:
                await self._get_table_structure(table_name)
    
    @handle_errors
    async def _get_table_structure(self, table_name: str) -> str:
        """
        从数据库中获取表结构
        
        Args:
            table_name: 表名
            
        Returns:
            str: 表结构信息
        """
        if table_name in self.table_schema_cache:
            return self.table_schema_cache[table_name]
        
        try:
            context = ToolContext(
                user=User(id="system", email="system@vanna.ai", group_memberships=['admin']),
                conversation_id="schema_loading",
                request_id="schema_request",
                agent_memory=self.agent_memory
            )
            
            query = f"SHOW CREATE TABLE {table_name}"
            result = await self.mysql_runner.run_sql(
                RunSqlToolArgs(sql=query),
                context
            )
            
            if not result.empty and 'Create Table' in result.columns:
                create_table_sql = result['Create Table'].iloc[0]
                table_schema = f"-- {self.mysql_runner.database}.`{table_name}` definition\n{create_table_sql}"
                
                # 缓存结果
                self.table_schema_cache[table_name] = table_schema
                logger.debug(f"获取表 {table_name} 结构成功")
                return table_schema
            
        except Exception as e:
            logger.error(f"获取表 {table_name} 结构时出错: {e}")
        
        return ""

    @handle_errors
    async def generate_business_contexts(self, table_names: Optional[list[str]] = None) -> list[str]:
        """
        动态生成业务表的建表语句上下文，并返回列表

        说明：
        - 优先使用传入的 `table_names`；未指定时使用 COMMON_TABLES；若为空则回退到全部表
        - 每条上下文包含完整的 `CREATE TABLE` 语句（字段、类型、约束等）
        """
        await self._ensure_initialized()

        names: list[str] = []
        if table_names:
            names = table_names
        elif COMMON_TABLES:
            names = COMMON_TABLES
        else:
            names = self.all_tables or []

        contexts: list[str] = []
        for name in names:
            schema = await self._get_table_structure(name)
            if schema:
                contexts.append(schema)

        # 若未能获取到指定表的结构，回退尝试数据库中的其他表（最多3个）
        if not contexts and self.all_tables:
            for name in self.all_tables:
                schema = await self._get_table_structure(name)
                if schema:
                    contexts.append(schema)
                if len(contexts) >= 3:
                    break

        return contexts
    
    @handle_errors
    async def search_documentation(self, query: str) -> str:
        """
        搜索相关的表结构信息，优先使用COMMON_TABLES中指定的表
        
        Args:
            query: 查询字符串
            
        Returns:
            str: 相关文档内容
        """
        await self._ensure_initialized()
        
        doc_parts = []
        
        # 确保COMMON_TABLES中指定的表结构都已加载
        for table_name in COMMON_TABLES:
            # 如果表结构尚未缓存，则获取它
            if table_name not in self.table_schema_cache:
                await self._get_table_structure(table_name)
            # 添加已缓存的表结构
            if table_name in self.table_schema_cache:
                doc_parts.append(self.table_schema_cache[table_name])
        
        # 如果没有COMMON_TABLES中的表结构，尝试获取其他可用表
        if not doc_parts and self.all_tables:
            # 优先获取数据库中存在的表
            for table_name in self.all_tables:
                if table_name not in self.table_schema_cache:
                    await self._get_table_structure(table_name)
                if table_name in self.table_schema_cache:
                    doc_parts.append(self.table_schema_cache[table_name])
                    # 最多获取3个表结构，避免信息过载
                    if len(doc_parts) >= 3:
                        break
        
        # 如果还是没有表结构，提供数据库信息
        if not doc_parts:
            doc_parts.append(f"-- 数据库名称: {self.mysql_runner.database}")
            if self.all_tables:
                doc_parts.append(f"-- 可用表: {', '.join(self.all_tables)}")
            # 提示用户配置的常用表
            doc_parts.append(f"-- 配置的常用表: {', '.join(COMMON_TABLES)}")
        
        return "\n\n".join(doc_parts)

    @handle_errors
    async def enhance_system_prompt(
        self,
        system_prompt: str,
        user_message: str,
        user: User
    ) -> str:
        """
        增强系统提示词，添加表结构信息和查询指导
        
        Args:
            system_prompt: 原始系统提示词
            user_message: 用户消息
            user: 用户对象
            
        Returns:
            str: 增强后的系统提示词
        """
        relevant_docs = await self.search_documentation(user_message)
        database_name = self.mysql_runner.database
        
        # 添加表结构信息
        system_prompt += f"\n\n相关表结构信息:\n{relevant_docs}"
        system_prompt += "\n\n重要提示："
        system_prompt += f"\n1. 数据库名为 {database_name}，请使用完整表名格式：{database_name}.表名"
        system_prompt += "\n2. 请严格按照上述表结构信息中定义的表名使用"
        system_prompt += "\n3. 请严格按照表结构中的字段名生成SQL查询"
        system_prompt += "\n4. 根据查询需求选择合适的表，可能需要查询多个表并使用JOIN操作"
        
        # 添加对数量查询的明确指导
        if any(keyword in user_message.lower() for keyword in ['多少', '数量', '总数', 'count']):
            system_prompt += "\n\n对于数量查询的特殊指导："
            system_prompt += "\n- 当用户询问数量相关问题时，请生成简单的COUNT查询"
            system_prompt += f"\n- 正确示例: SELECT COUNT(*) AS total_rows FROM {database_name}.表名"
            system_prompt += "\n- 不要添加额外的WHERE条件，除非用户明确指定了过滤条件"
            system_prompt += "\n- 不要在SELECT子句中包含除COUNT外的其他字段，除非用户明确要求"
        
        system_prompt += "\n\n请根据用户的问题和提供的表结构信息，生成准确、高效的SQL查询。"
        return system_prompt

    @handle_errors
    async def enhance_user_messages(
        self,
        messages: list[LlmMessage],
        user: User
    ) -> list[LlmMessage]:
        """
        增强用户消息，添加用户上下文信息
        
        Args:
            messages: 原始消息列表
            user: 用户对象
            
        Returns:
            list[LlmMessage]: 增强后的消息列表
        """
        messages.append(LlmMessage(
            role="system", 
            content=f"[系统上下文信息：当前登录用户为 {user.email}，所属组为 {', '.join(user.group_memberships)}。此信息仅用于身份验证和权限控制，不是用户查询条件。]"
        ))
        return messages


async def create_production_agent() -> Agent:
    """
    创建生产环境代理实例
    
    Returns:
        Agent: 配置好的代理实例
    """
    # 加载配置
    db_config = DatabaseConfig.from_env()
    llm_config = LLMConfig.from_env()
    server_config = ServerConfig.from_env()
    
    logger.info("正在初始化生产环境代理...")
    
    # 初始化LLM服务
    llm = OllamaLlmService(
        model=llm_config.model,
        host=llm_config.host
    )
    logger.info(f"LLM服务已初始化: {llm_config.model} @ {llm_config.host}")
    
    # 初始化数据库工具
    db_tool = ProductionEditableRunSqlTool(
        sql_runner=MySQLRunner(
            host=db_config.host,
            database=db_config.database,
            user=db_config.user,
            password=db_config.password,
            port=db_config.port
        )
    )
    logger.info(f"数据库连接已初始化: {db_config.user}@{db_config.host}:{db_config.port}/{db_config.database}")
    
    # 初始化代理内存
    agent_memory = ChromaAgentMemory(
        collection_name="vanna_memory",
        persist_directory="./chroma_db"
    )
    logger.info("代理内存已初始化")
    
    # 初始化用户解析器
    user_resolver = ProductionUserResolver()
    logger.info("用户解析器已初始化")
    
    # 注册工具
    tools = ToolRegistry()
    tools.register_local_tool(db_tool, access_groups=['admin', 'user'])
    tools.register_local_tool(SaveQuestionToolArgsTool(), access_groups=['admin'])
    tools.register_local_tool(SearchSavedCorrectToolUsesTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(SaveTextMemoryTool(), access_groups=['admin', 'user'])
    tools.register_local_tool(VisualizeDataTool(), access_groups=['admin', 'user'])
    logger.info("工具注册完成")
    
    # 配置代理参数
    config = AgentConfig(
        max_tool_iterations=server_config.max_tool_iterations,
        stream_mode=True,
        temperature=server_config.temperature,
        include_tool_names=True,
        auto_save_context=True,
        max_tokens=None
    )
    
    # 创建代理
    agent = Agent(
        llm_service=llm,
        tool_registry=tools,
        user_resolver=user_resolver,
        agent_memory=agent_memory,
        llm_context_enhancer=ProductionLlmContextEnhancer(
            mysql_runner=db_tool.sql_runner, 
            agent_memory=agent_memory
        ),
        config=config
    )
    
    # 动态生成并缓存业务表建表语句到 business_contexts（参考 tests/test5.py 的文档入库方式）
    try:
        enhancer: ProductionLlmContextEnhancer = agent.llm_context_enhancer  # type: ignore
        contexts = await enhancer.generate_business_contexts(COMMON_TABLES)
        # 保存到模块级变量，供其它测试/脚本使用
        global business_contexts
        business_contexts = contexts
        logger.info(f"业务建表语句已生成并缓存，共 {len(business_contexts)} 条")
        
        # 将建表语句以文本形式写入代理内存（与 tests/test5.py 风格一致）
        admin_user = User(id="admin", email="admin@example.com", group_memberships=['admin'])
        ctx = ToolContext(
            user=admin_user,
            conversation_id="business_schema_bootstrap",
            request_id="bootstrap_request",
            agent_memory=agent.agent_memory
        )
        for text in business_contexts:
            await agent.agent_memory.save_text_memory(content=text, context=ctx)
            # 缓存数据：将建表语句写入业务上下文变量
            logger.info(f"已缓存建表语句: {text[:100]}...")
        logger.info("业务建表语句已写入代理内存")
    except Exception as e:
        logger.warning(f"生成或写入业务建表语句失败: {e}")
    
    logger.info("生产环境代理初始化完成")
    return agent


async def create_app():
    """
    创建应用实例，用于uvicorn启动
    
    Returns:
        FastAPI应用实例
    """
    # 创建代理
    agent = await create_production_agent()
    
    # 创建并返回FastAPI应用
    server = VannaFastAPIServer(agent=agent)
    app = server.create_app()
    logger.info("FastAPI应用实例已创建，可用于uvicorn启动")
    
    return app


def create_app_sync():
    """
    同步包装函数，用于uvicorn工厂模式
    
    Returns:
        FastAPI应用实例
    """
    # 使用nest_asyncio运行异步函数
    import asyncio
    try:
        loop = asyncio.get_event_loop()
    except RuntimeError:
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
    
    app = loop.run_until_complete(create_app())
    return app


async def main():
    """主函数，启动生产环境服务器"""
    try:
        # 获取服务器配置
        server_config = ServerConfig.from_env()
        
        # 创建应用实例
        app = await create_app()
        
        # 创建并启动服务器
        logger.info(f"正在启动生产环境服务器，端口: {server_config.port}")
        
        # 获取uvicorn配置
        import uvicorn.config
        
        # 创建uvicorn配置
        config = uvicorn.config.Config(
            app=app,
            host="0.0.0.0",
            port=server_config.port,
            log_level="info"
        )
        
        # 创建服务器实例
        uvicorn_server = uvicorn.Server(config)
        
        # 启动服务器
        await uvicorn_server.serve()
        
    except Exception as e:
        logger.error(f"启动服务器时发生错误: {e}", exc_info=True)
        raise

if __name__ == "__main__":
    # 运行主函数
    asyncio.run(main())
